from fastapi import APIRouter, HTTPException, Body, Depends
from api.common import DocumentCreate, DocumentResponse, get_document_manager, get_vector_store
from pydantic import BaseModel
from db.document_manager import DocumentManager
from db.model_loader import model_loader

router = APIRouter(prefix="/documents", tags=["文档"])


@router.get("/", response_model=list[DocumentResponse])
async def get_all_documents(doc_manager=Depends(get_document_manager)):
    """获取所有文档"""
    documents = doc_manager.get_documents_by_category(None)  # 传入None获取所有文档

    # 按更新时间排序（降序），只返回前5篇
    documents_sorted = sorted(documents, key=lambda x: x[5], reverse=True)[:5]

    return [
        {
            "id": doc[0],
            "content": doc[1],
            "title": doc[2],
            "category_ids": doc[3],
            "created_at": str(doc[4]),
            "updated_at": str(doc[5]),
        }
        for doc in documents_sorted
    ]


@router.get("/count")
async def get_document_count(doc_manager: DocumentManager = Depends(get_document_manager)):
    """获取文档总数"""
    try:
        total_count = doc_manager.get_total_documents()
        return {"total_documents": total_count}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取文档总数失败: {str(e)}")


@router.post("/", response_model=DocumentResponse)
async def create_document(document: DocumentCreate, vector_store=Depends(get_vector_store)):
    """创建新的文档并存储其向量嵌入"""
    try:
        doc_id = vector_store.store_document(document.content, document.category_ids, document.title)
        # 这里简化了返回值，实际项目中应查询数据库获取完整信息
        return {
            "id": doc_id,
            "content": document.content,
            "category_ids": document.category_ids,
            "title": document.title,
            "created_at": "刚刚创建",
            "updated_at": "刚刚创建",
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"创建文档失败: {str(e)}")


@router.get("/category/{category_id}", response_model=list[DocumentResponse])
async def get_documents_by_category(category_id: int, doc_manager=Depends(get_document_manager)):
    """获取指定分类下的所有文档"""
    documents = doc_manager.get_documents_by_category(category_id)
    return [
        {"id": doc[0], "content": doc[1], "title": doc[2], "category_ids": doc[3], "created_at": str(doc[4]), "updated_at": str(doc[5])}
        for doc in documents
    ]


@router.get("/{document_id}", response_model=DocumentResponse)
async def get_document(document_id: int, doc_manager=Depends(get_document_manager)):
    """通过ID获取文档"""
    document = doc_manager.get_document_by_id(document_id)
    if not document:
        raise HTTPException(status_code=404, detail="文档不存在")
    return {
        "id": document[0],
        "content": document[1],
        "category_ids": document[2],
        "title": document[3],
        "created_at": str(document[4]),
        "updated_at": str(document[5]),
    }


@router.put("/{document_id}/categories")
async def update_document_categories_endpoint(
    document_id: int, category_ids: list[int] = Body(..., embed=True), doc_manager=Depends(get_document_manager)
):
    """更新文档的分类（支持多个分类）"""
    try:
        # 检查文档是否存在
        document = doc_manager.get_document_by_id(document_id)
        if not document:
            raise HTTPException(status_code=404, detail="文档不存在")

        # 检查所有分类是否存在
        for cat_id in category_ids:
            category = doc_manager.get_category_by_id(cat_id)
            if not category:
                raise HTTPException(status_code=404, detail=f"分类ID {cat_id} 不存在")

        # 更新分类
        doc_manager.update_document_categories(document_id, category_ids)
        return {"message": "文档分类更新成功", "category_ids": category_ids}
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"更新文档分类失败: {str(e)}")


@router.put("/{document_id}", response_model=DocumentResponse)
async def update_document_endpoint(document_id: int, document: DocumentCreate, doc_manager=Depends(get_document_manager)):
    """更新文档内容、标题和分类"""
    try:
        # 检查文档是否存在
        existing_document = doc_manager.get_document_by_id(document_id)
        if not existing_document:
            raise HTTPException(status_code=404, detail="文档不存在")

        # 检查所有分类是否存在
        if document.category_ids:
            for cat_id in document.category_ids:
                category = doc_manager.get_category_by_id(cat_id)
                if not category:
                    raise HTTPException(status_code=404, detail=f"分类ID {cat_id} 不存在")

        # 更新文档
        success = doc_manager.update_document(document_id, title=document.title, content=document.content, category_ids=document.category_ids)

        if not success:
            raise HTTPException(status_code=404, detail="文档不存在")

        # 获取更新后的文档信息
        updated_document = doc_manager.get_document_by_id(document_id)
        if not updated_document:
            raise HTTPException(status_code=500, detail="更新文档后获取文档信息失败")

        return {
            "id": updated_document[0],
            "content": updated_document[1],
            "category_ids": updated_document[2],
            "title": updated_document[3],
            "created_at": str(updated_document[4]),
            "updated_at": str(updated_document[5]),
        }
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"更新文档失败: {str(e)}")


class SmartCategoryRequest(BaseModel):
    content: str

class FormatMarkdownRequest(BaseModel):
    content: str


@router.post("/format-markdown")
async def format_markdown_document(request: FormatMarkdownRequest):
    """将文档内容格式化为Markdown"""
    try:
        # 构建提示词，请求将内容格式化为Markdown
        prompt = (
            "你是一个专业的Markdown格式编辑助手。请将以下文档内容转换为标准的Markdown格式。\n"
            "要求：\n"
            "1. 合理使用标题（#、##等）、列表（有序和无序）、粗体、斜体等Markdown元素\n"
            "2. 保持原文的语义和结构，不要改动任何的原始内容\n"
            "3. 添加适当的段落分隔\n"
            "4. 如有代码块，请使用```包裹\n"
            "5. 如有表格，请使用Markdown表格语法\n"
            "6. 不要在返回内容中添加任何额外的解释性文字\n\n"
            f"文档内容：\n{request.content}"
        )
        print(prompt)
        print("=" * 50)

        # 调用LLM生成Markdown格式的内容
        try:
            # 统一通过model_loader获取LLM响应，model_loader会根据API_TYPE选择合适的方式
            # 增加max_new_tokens参数，确保能返回完整的格式化内容
            response = model_loader.generate_llm_response(prompt, max_new_tokens=2000)
        except Exception as e:
            print(f"格式化失败: {e}")
            raise HTTPException(status_code=500, detail=f"格式化文档失败: {str(e)}")

        print("格式化结果:")
        print(response)
        
        # 返回格式化后的Markdown内容
        return {"formatted_content": response}
    except Exception as e:
        print(f"格式化文档失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"格式化文档失败: {str(e)}")

@router.post("/smart-categorize")
async def smart_categorize_document(request: SmartCategoryRequest, doc_manager=Depends(get_document_manager)):
    """智能分类文档，使用LLM预测文档所属分类"""
    try:
        # 获取所有已存在的分类
        existing_categories = doc_manager.get_all_categories()
        category_names = [cat[1] for cat in existing_categories]

        # 构建提示词
        prompt = (
            f"我现在需要给一篇文档进行智能分类，现有分类有：{','.join(category_names)}。\n"
            f"请你根据文档内容，输出5个分类，每个类别不超过6个字符，也不少于一个字符，\n"
            f"请在回答的开头就直接输出分类，其它什么都不要输出，不要输出任何的思考过程，格式为：#分类1；分类2；分类3；分类4；分类5#\n"
            f"优先从已存在的分类中选择，但要求至少有2个新增的分类，如果符合要求的分类不够，则新增符合文档的分类，\n"
            "所有的分类用中文分号隔开。第一个分类的前面加上井号，最后一个分类的后面加上井号。\n\n"
            f"文档内容：{request.content[:2000]}"  # 限制输入长度，避免超出模型上下文
        )
        print(prompt)
        print("=" * 50)

        # 调用LLM生成分类结果
        try:
            # 统一通过model_loader获取LLM响应，model_loader会根据API_TYPE选择合适的方式
            response = model_loader.generate_llm_response(prompt)
        except Exception as e:
            print(f"分类失败: {e}")
            # 如果失败，返回一个默认的分类
            # return [{"name": "其他", "confidence": 1.0}]

        # 解析生成的分类
        # 查找井号作为结束标记
        print(response)
        if "#" in response:
            categories_text = response.split("#")[1].strip()
        else:
            categories_text = response.strip()

        # 分割分类
        categories = [cat.strip() for cat in categories_text.split("；") if cat.strip()]

        # 处理可能使用英文分号的情况
        if len(categories) <= 1 and ";" in categories_text:
            categories = [cat.strip() for cat in categories_text.split(";") if cat.strip()]

        # 限制最多5个分类
        categories = categories[:5]

        # 检查每个分类是否已存在
        result = []
        for cat_name in categories:
            # 检查是否超过6个字符
            if len(cat_name) > 6:
                continue

            # 查找是否为已存在的分类
            existing_cat = next((cat for cat in existing_categories if cat[1] == cat_name), None)
            if existing_cat:
                result.append({"id": existing_cat[0], "name": existing_cat[1], "is_existing": True})
            else:
                result.append({"id": None, "name": cat_name, "is_existing": False})

        # 如果生成的分类不足5个，可以根据需要补充
        if len(result) < 5:
            # 这里可以添加一些逻辑来补充分类，但为了简单起见，我们保持原样
            pass

        return {"categories": result}
    except Exception as e:
        print(f"智能分类失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"智能分类失败: {str(e)}")


@router.delete("/{document_id}")
async def delete_document_endpoint(document_id: int, doc_manager=Depends(get_document_manager)):
    """删除文档"""
    try:
        # 检查文档是否存在
        document = doc_manager.get_document_by_id(document_id)
        if not document:
            raise HTTPException(status_code=404, detail="文档不存在")

        # 删除文档
        success = doc_manager.delete_document(document_id)

        if not success:
            raise HTTPException(status_code=404, detail="文档不存在")

        return {"message": "文档删除成功"}
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"删除文档失败: {str(e)}")
